#include "../../include/sm2_params.h"
#include <openssl/ecdh.h>
#include <openssl/ec.h>
#include <openssl/crypto.h>
#include <openssl/obj_mac.h> // 添加此头文件以获取NID定义
#include <string.h>
#include <openssl/evp.h> // Added for EVP_PKEY functions

// 如果系统没有提供NID_sm2，提供一个自定义定义
#ifndef NID_sm2
#define NID_sm2 1172 // 根据OpenSSL定义的SM2曲线ID
#endif

// SM3 KDF (密钥派生函数)实现
static int SM3_KDF(const unsigned char *Z, size_t zlen,
                   unsigned char *K, size_t klen)
{
    // 简化实现，实际项目中应完整实现SM3 KDF
    // 这里仅作为示例
    if (K && klen > 0)
    {
        memset(K, 0x42, klen); // 临时占位
        return 1;
    }
    return 0;
}

int ecdh_sm2_derive_key(EC_KEY *local_key, const EC_POINT *peer_public,
                        unsigned char *skey, size_t skey_len)
{
    const EC_GROUP *group = EC_KEY_get0_group(local_key);
    if (!group)
    {
        return 0;
    }

    EVP_PKEY *local_evp_key = NULL;
    EVP_PKEY *peer_evp_key = NULL;
    EC_KEY *peer_ec_key = NULL;
    EVP_PKEY_CTX *derive_ctx = NULL;
    unsigned char *derived_secret = NULL;
    size_t derived_secret_len = 0;
    int ret = 0;

    // Wrap local EC_KEY in EVP_PKEY
    local_evp_key = EVP_PKEY_new();
    if (!local_evp_key || EVP_PKEY_set1_EC_KEY(local_evp_key, local_key) != 1)
    {
        goto err;
    }

    // Create peer EC_KEY from public point and group, then wrap in EVP_PKEY
    peer_ec_key = EC_KEY_new();
    if (!peer_ec_key || EC_KEY_set_group(peer_ec_key, group) != 1 ||
        EC_KEY_set_public_key(peer_ec_key, peer_public) != 1)
    {
        goto err;
    }
    peer_evp_key = EVP_PKEY_new();
    if (!peer_evp_key || EVP_PKEY_set1_EC_KEY(peer_evp_key, peer_ec_key) != 1)
    {
        goto err;
    }

    // Create context for key derivation
    derive_ctx = EVP_PKEY_CTX_new(local_evp_key, NULL);
    if (!derive_ctx)
    {
        goto err;
    }

    // Initialize derivation
    if (EVP_PKEY_derive_init(derive_ctx) <= 0)
    {
        goto err;
    }

    // Set peer public key
    if (EVP_PKEY_derive_set_peer(derive_ctx, peer_evp_key) <= 0)
    {
        goto err;
    }

    // Determine buffer length for derived secret
    if (EVP_PKEY_derive(derive_ctx, NULL, &derived_secret_len) <= 0)
    {
        goto err;
    }

    derived_secret = OPENSSL_malloc(derived_secret_len);
    if (!derived_secret)
    {
        goto err;
    }

    // Derive the shared secret
    if (EVP_PKEY_derive(derive_ctx, derived_secret, &derived_secret_len) <= 0)
    {
        goto err;
    }

    // Use the derived secret with your KDF
    // Ensure skey_len is appropriate for SM3_KDF output
    if (!SM3_KDF(derived_secret, derived_secret_len, skey, skey_len))
    {
        goto err;
    }

    ret = 1; // Success

err:
    EVP_PKEY_free(local_evp_key);
    EC_KEY_free(peer_ec_key); // Free the temporary peer EC_KEY
    EVP_PKEY_free(peer_evp_key);
    EVP_PKEY_CTX_free(derive_ctx);
    OPENSSL_free(derived_secret);

    return ret;
}
